import jieba
import numpy as np
from rank_bm25 import BM25Okapi
from scipy.special import softmax

from common.logger import logger
from mysql_qa.mysql_client import MySQLClient
from mysql_qa.redis_client import RedisClient


def clean(lines):
    if not lines:
        return None
    result = []
    for line in lines:
        question = preprocess(line[1])
        if not question:
            continue
        result.append((line[0], question))
    return result


def preprocess(content):
    try:
        return jieba.lcut(content.strip().lower())
    except Exception as e:
        logger.error(e)


class BM25Search:
    """
    redis存储命中的高相似问答对
    mysql存储全量问答对
    """

    def __init__(self, redis_client: RedisClient, mysql_client: MySQLClient):
        self.redis_client = redis_client
        self.mysql_client = mysql_client
        self.bm25 = None
        self.questions = None
        self.question_ids = None
        self._load_data()

    def _load_data(self):
        """
        查询全部 qa，初始化 BM25
        """
        lines = self.mysql_client.fetch_all_questions()
        lines = clean(lines)
        if not lines:
            raise Exception("未加载到问题")

        self.question_ids = [line[0] for line in lines]
        self.questions = [line[1] for line in lines]
        self.bm25 = BM25Okapi(self.questions)
        logger.info("BM25 模型初始化完成")

    def search(self, query: str, threshold=0.85):
        query = preprocess(query)
        if not query:
            logger.error("无效查询", query)
            return None
        # 计算BM25
        scores = self.bm25.get_scores(query)
        softmax_scores: np.ndarray = softmax(scores)
        best_idx = softmax_scores.argmax()
        softmax_best_score = softmax_scores[best_idx]
        logger.info(
            f'query={query}\nbest_idx={best_idx}, best_score={scores[best_idx]}, softmax_best_score={softmax_best_score}')
        # 满足阈值，依次从 cache，db 检索答案返回
        if softmax_best_score >= threshold:
            question_id = self.question_ids[best_idx]
            logger.info(f'满足阈值, question_id={question_id}')
            answer = self.redis_client.get_with('answer', question_id)
            if not answer:
                logger.info(f'not hit redis, query DB')
                answer = self.mysql_client.get_answer(question_id)
                self.redis_client.set_with('answer', question_id, answer)
            else:
                logger.info(f'命中redis')
            return answer
        return None
